package org.ollide.java2smali

import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.compiler.CompilerPaths
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.module.Module
import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.*
import com.intellij.openapi.vfs.JarFileSystem
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.psi.PsiClassOwner
import com.intellij.psi.PsiManager
import com.intellij.task.ProjectTaskManager
import com.intellij.util.SmartList
import com.intellij.util.containers.OrderedSet
import com.intellij.util.io.URLUtil
import org.beansoft.asm.idea.plugin.view.SmaliView
import java.io.File
import java.io.IOException
import java.nio.file.Files
import java.nio.file.Paths

class DexCompiler(private val vFile: VirtualFile, private val project: Project, private var module: Module, private val needBuild:Boolean = true) {

    fun run() {
        if(needBuild) {
            buildModule {
                onProjectBuildComplete()
            }
        } else {
            onProjectBuildComplete()
        }

    }

    /**
     * To create a dex or smali file from the virtual file, we ne a compiled .class file
     * of the given virtual file to be present.
     *
     * Project structures and builds vary a lot (Android, Java, Kotlin, directory structure),
     * so instead of using the CompileManager, we trigger a general build task.
     */
    private fun buildModule(callback: () -> Unit) {
        val projectTaskManager = ProjectTaskManager.getInstance(project)
        val buildTask = projectTaskManager.createModulesBuildTask(module, true, true, true)

        val supportProjectTaskManager = SupportProjectTaskManager(projectTaskManager)
        supportProjectTaskManager.run(buildTask).onSuccess {
            if (it.hasErrors()) {
                LOG.warn("Module build failed, aborting dex/smali build.")
            } else {
                callback()
            }
        }
    }

    private fun onProjectBuildComplete() {
        val psiClassOwner = PsiManager.getInstance(project).findFile(vFile) as PsiClassOwner
        var classVirtualFile:VirtualFile = vFile
        val isClass = "class".equals(vFile.extension, ignoreCase = true)
        var isInJar = vFile.fileSystem is JarFileSystem // jar:///
        val virtualFileManager = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)

        if(isInJar && !isClass) {
            setSmaliCode("Ignored: only class file in jar file supported")
            return
        }
        // 找不到输出目录时退出
        var dexFileOutputDirectory:VirtualFile
        if(isInJar) {
            val tempDirectory = Files.createTempDirectory("android-decompiler-")
            dexFileOutputDirectory = LocalFileSystem.getInstance().findFileByIoFile(tempDirectory.toFile())!!
        } else {
            try {
                dexFileOutputDirectory = getFileOutputDirectory(psiClassOwner)
            } catch (e: Exception) {
                try {
                    dexFileOutputDirectory = vFile.parent//  getFileOutputDirectorySimple(file)
                } catch (e: Exception) {
                    return
                }
            }
        }

        // Copy jar class file content to temp output folder
        if(isInJar && isClass) {
            val classFile = Paths.get(dexFileOutputDirectory.path, vFile.nameWithoutExtension + CLASS_EXTENSION)
            try {
                Files.newOutputStream(classFile).use { out ->
                    out.write(vFile.contentsToByteArray())
                }
            } catch (e: java.lang.Exception) {
                throw IOException("Failed to save temp file", e)
            }

            classVirtualFile = LocalFileSystem.getInstance().findFileByIoFile(classFile.toFile())!!

            println(classVirtualFile)
//            virtualFileManager.refresh(false)
        }

        dexFileOutputDirectory.refresh(false, false)

        val fileName = vFile.nameWithoutExtension
        val dexFilePath = Paths.get(dexFileOutputDirectory.path, fileName + DEX_EXTENSION).toString()

        // CLASS -> DEX
        var targetClassFiles = getClassFiles(dexFileOutputDirectory, fileName)
        compileDexFile(targetClassFiles, dexFilePath)

        // DEX -> SMALI
//        val outputDir = getSourceRootFile().path // 输出smail到 src 目录

        var outputDir:String
        if(isInJar) {
            outputDir = dexFileOutputDirectory.path
        } else {
            try {
                outputDir = if(isClass) getRootOutputDirectorySimple(psiClassOwner).path else {
                    getRootOutputDirectory().path // 输出目录可能为空
                }
            } catch (e: Exception) {
                try {
                    outputDir = getRootOutputDirectorySimple(psiClassOwner).path
                } catch (e: Exception) {
                    setSmaliCode("Error: " + e.message)
                    return
                }
            }
        }

        WriteCommandAction.runWriteCommandAction(project) {
            try {
                Dex2SmaliHelper.disassembleDexFile(dexFilePath, outputDir)
            } catch (e: Exception) {
                val smaliView = SmaliView.getInstance(project)
                WriteCommandAction.runWriteCommandAction(project) {
                    smaliView.setCode(
                            null,  "Failed to decompile smali:" + e.message)
                }
                return@runWriteCommandAction
            }

            // we've created the smali file(s) in our source file's directory
            // refresh directory synchronously and access children to let IDEA detect the file(s)
//            val parent = vFile.parent
//            parent.refresh(false, false)
//            parent.children
            dexFileOutputDirectory.refresh(false, true)
            // get a VirtualFile by the IO path
//            val smaliPath = vFile.path.substringBeforeLast('.') + SMALI_EXTENSION
            var smaliPath = (vFile.path.substringBeforeLast('.') + SMALI_EXTENSION)
            if (getSourceRootFile() != null) {
                smaliPath = smaliPath.replaceFirst(getSourceRootFile()!!.path, outputDir)
            }

            if(isInJar) {
                val pkg = psiClassOwner.packageName.replace('.', '/')
                smaliPath = Paths.get(dexFileOutputDirectory.path, pkg, fileName + SMALI_EXTENSION).toString()
            }

            var virtualSmaliFile:VirtualFile = LocalFileSystem.getInstance().findFileByIoFile(File(smaliPath)) ?: return@runWriteCommandAction

            // use the VirtualFile to show the smali file in IDEA editor
//        val openFileDescriptor = OpenFileDescriptor(project, virtualDexFile)
//        openFileDescriptor.navigate(true)
            // use the VirtualFile to show the smali file in IDEA editor
//        val openFileDescriptor = OpenFileDescriptor(project, virtualDexFile)
//        openFileDescriptor.navigate(true)
            showSmali(virtualSmaliFile)
        }

    }

    private fun showSmali(smaliFile : VirtualFile) {
        val smaliView = SmaliView.getInstance(project)
        WriteCommandAction.runWriteCommandAction(project) {
            smaliView.setCode(
                smaliFile,  String(smaliFile.contentsToByteArray()).replace("\r\n", "\n")
            )
        }
    }

    private fun setSmaliCode(code : String) {
        val smaliView = SmaliView.getInstance(project)
        WriteCommandAction.runWriteCommandAction(project) {
            smaliView.setCode(
                    null,  code
            )
        }
    }

    private fun getClassFiles(fileOutputDirectory: VirtualFile, fileName: String): Array<String> {
        // Class 文件返回自身
        if ("class".equals(fileOutputDirectory.extension, ignoreCase = true)) {
            return arrayOf(fileOutputDirectory.path)
        }

        val children = fileOutputDirectory.children ?: arrayOf()
        return children.filter {
            val baseName = it.nameWithoutExtension
            (baseName == fileName || baseName.startsWith("$fileName$")) && it.extension == CLASS
        }.map {
            it.path
        }.toTypedArray()
    }

    private fun getFileOutputDirectory(file: PsiClassOwner): VirtualFile {
        // determine whether this is a production or test file
        val isProduction = module.getModuleScope(false).contains(vFile)

        val pkg = file.packageName.replace('.', File.separatorChar)

        // find the general output directory of the file's module (target, app/build/intermediates/javac/$variant/classes, ...)
        val possibleOutputDirectories = findModuleOutputDirectories(isProduction)
        LOG.debug("Possible output directories: ", possibleOutputDirectories.joinToString(","))

        val virtualFileManager = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)

        val fileOutputDirectory = possibleOutputDirectories
            .asSequence()
            .map {
                val classFile = vFile.nameWithoutExtension + CLASS_EXTENSION
                val path = Paths.get(it, pkg, classFile).toString()
                virtualFileManager.refreshAndFindFileByPath(path)?.parent
            }
            .firstOrNull { it != null }

        LOG.debug("Found output directory: $fileOutputDirectory")
        return fileOutputDirectory ?: throw IllegalStateException("Output directory not found")
    }

    private fun getFileOutputDirectorySimple(file: PsiClassOwner): VirtualFile {
        val path = vFile.path
        val pkg = file.packageName.replace('.', '/')

        val classFile = pkg + "/" + vFile.name

        val outputDir = path.replace(classFile, "")

        val virtualFileManager = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)

        val fileOutputDirectory = virtualFileManager.refreshAndFindFileByPath(outputDir)

        LOG.debug("Found output directory: $fileOutputDirectory")
        return fileOutputDirectory ?: throw IllegalStateException("Output directory not found")
    }

    private fun getRootOutputDirectorySimple(file: PsiClassOwner): VirtualFile {
        val path = vFile.path
        val pkg = file.packageName.replace('.', '/')

        val classFile = pkg + "/" + vFile.name

        val outputDir = path.replace(classFile, "")

        val virtualFileManager = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)

        val fileOutputDirectory = virtualFileManager.refreshAndFindFileByPath(outputDir)

        LOG.debug("Found output directory: $fileOutputDirectory")
        return fileOutputDirectory ?: throw Exception("Output directory not found")
    }

    // 获取项目输出根目录
    private fun getRootOutputDirectory(): VirtualFile {
        // determine whether this is a production or test file
        val isProduction = module.getModuleScope(false).contains(vFile)

        // find the general output directory of the file's module (target, app/build/intermediates/javac/$variant/classes, ...)
        val possibleOutputDirectories = findModuleOutputDirectories(isProduction)
        LOG.debug("Possible output directories: ", possibleOutputDirectories.joinToString(","))

        val virtualFileManager = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)

        val fileOutputDirectory = possibleOutputDirectories
            .asSequence()
            .map {
                virtualFileManager.refreshAndFindFileByPath(it)
            }
            .firstOrNull { it != null }

        LOG.debug("Found output directory: $fileOutputDirectory")
        return fileOutputDirectory ?: throw IllegalStateException("Output directory not found")
    }

    /**
     * @see <a href="https://github.com/JetBrains/intellij-community/blob/master/java/compiler/openapi/src/com/intellij/openapi/compiler/CompilerPaths.java">
     *     intellij-community/CompilerPaths.java</a>
     *     method: public static String @NotNull [] getOutputPaths(Module @NotNull [] modules)
     */
    private fun findModuleOutputDirectories(production: Boolean): OrderedSet<String> {
        /*
        val outputPaths: MutableList<String> = mutableListOf()

        val compilerExtension = CompilerModuleExtension.getInstance(module)
//        if (production) {
//            compilerExtension?.compilerOutputPath?.path?.let { outputPaths.add(it) }
//        } else {
//            compilerExtension?.compilerOutputPathForTests?.path?.let { outputPaths.add(it) }
//        }

        compilerExtension?.compilerOutputPath?.path?.let { outputPaths.add(it) }
        compilerExtension?.compilerOutputPathForTests?.path?.let { outputPaths.add(it) }

        val moduleRootManager = ModuleRootManager.getInstance(module)
        for (handlerFactory in OrderEnumerationHandler.EP_NAME.extensions) {
            if (handlerFactory.isApplicable(module)) {
                val handler = handlerFactory.createHandler(module)
                val outputUrls: List<String> = SmartList()
                handler.addCustomModuleRoots(OrderRootType.CLASSES, moduleRootManager, outputUrls, production, !production)
                for (outputUrl in outputUrls) {
                    outputPaths.add(VirtualFileManager.extractPath(outputUrl).replace('/', File.separatorChar))
                }
            }
        }
*/
        val outputPathStrings = CompilerPaths.getOutputPaths(arrayOf(module));

        return OrderedSet(outputPathStrings.asList())
    }

    private fun compileDexFile(compiledPaths: Array<String>, dexFile: String) {
        try {
            Class2DexHelper.dexClassFile(compiledPaths, dexFile)
        } catch (e: Throwable) {
            e.printStackTrace()
            return
        }
    }

    private fun getSourceRootFile(): VirtualFile? {
        try {
            return ProjectRootManager.getInstance(project).fileIndex.getSourceRootForFile(vFile) as VirtualFile
        } catch (e: Exception) {
            return null
        }
    }

    companion object {

        private val LOG = Logger.getInstance(DexCompiler::class.java)

        const val CLASS_EXTENSION = ".class"
        const val DEX_EXTENSION = ".dex"
        const val SMALI_EXTENSION = ".smali"
        const val CLASS = "class"
    }
}
